#include <mcp/server/initialization.h>
#include <mcp/error.h>
#include <mcp/json/value.h>
#include <mcp/types.h>
#include <sstream>
#include <algorithm>

namespace mcp {
namespace server {

// ServerCapabilities implementation
ServerCapabilities::ServerCapabilities(
    bool tools_enabled,
    bool resources_enabled,
    bool prompts_enabled,
    bool logging_enabled,
    bool progress_enabled,
    std::optional<ExperimentalCapabilities> experimental)
    : tools(tools_enabled
        ? ToolsCapabilities{.list_changed = true}
        : std::optional<ToolsCapabilities>())
    , resources(resources_enabled
        ? ResourcesCapabilities{.subscribe = true, .list_changed = true}
        : std::optional<ResourcesCapabilities>())
    , prompts(prompts_enabled
        ? PromptsCapabilities{.list_changed = true}
        : std::optional<PromptsCapabilities>())
    , logging(logging_enabled
        ? LoggingCapabilities{.level = LoggingLevel::Debug}
        : std::optional<LoggingCapabilities>())
    , progress(progress_enabled)
    , experimental(std::move(experimental)) {}

ServerCapabilities ServerCapabilities::create_default() {
    return ServerCapabilities{
        .tools_enabled = true,
        .resources_enabled = true,
        .prompts_enabled = true,
        .logging_enabled = true,
        .progress_enabled = true,
        .experimental = std::nullopt
    };
}

ServerCapabilities ServerCapabilities::create_minimal() {
    return ServerCapabilities{
        .tools_enabled = true,
        .resources_enabled = false,
        .prompts_enabled = false,
        .logging_enabled = false,
        .progress_enabled = false,
        .experimental = std::nullopt
    };
}

json::Value ServerCapabilities::to_json() const {
    json::Value caps = json::Value::object({});

    if (tools) {
        caps["tools"] = json::Value::object({
            {"listChanged", tools->list_changed}
        });
    }

    if (resources) {
        caps["resources"] = json::Value::object({
            {"subscribe", resources->subscribe},
            {"listChanged", resources->list_changed}
        });
    }

    if (prompts) {
        caps["prompts"] = json::Value::object({
            {"listChanged", prompts->list_changed}
        });
    }

    if (logging) {
        std::string level_str;
        switch (logging->level) {
            case LoggingLevel::Debug:
                level_str = "debug";
                break;
            case LoggingLevel::Info:
                level_str = "info";
                break;
            case LoggingLevel::Warning:
                level_str = "warning";
                break;
            case LoggingLevel::Error:
                level_str = "error";
                break;
            default:
                level_str = "info";
                break;
        }
        caps["logging"] = json::Value::object({
            {"level", level_str}
        });
    }

    if (progress) {
        caps["progress"] = true;
    }

    if (experimental) {
        caps["experimental"] = json::Value::object({});
        for (const auto& [key, value] : *experimental) {
            caps["experimental"][key] = value;
        }
    }

    return caps;
}

// InitializationOptions implementation
InitializationOptions::InitializationOptions(
    std::string server_name,
    std::string server_version,
    ServerCapabilities capabilities,
    std::optional<std::string> instructions)
    : server_name(std::move(server_name))
    , server_version(std::move(server_version))
    , capabilities(std::move(capabilities))
    , instructions(std::move(instructions)) {}

json::Value InitializationOptions::to_json() const {
    json::Value result = json::Value::object({
        {"protocolVersion", "2025-06-18"},
        {"capabilities", capabilities.to_json()},
        {"serverInfo", json::Value::object({
            {"name", server_name},
            {"version", server_version}
        })}
    });

    if (instructions) {
        result["instructions"] = *instructions;
    }

    return result;
}

// InitializationResult implementation
InitializationResult::InitializationResult(
    std::string protocol_version,
    ServerCapabilities capabilities,
    types::Implementation server_info,
    std::optional<std::string> instructions)
    : protocol_version(std::move(protocol_version))
    , capabilities(std::move(capabilities))
    , server_info(std::move(server_info))
    , instructions(std::move(instructions)) {}

json::Value InitializationResult::to_json() const {
    json::Value result = json::Value::object({
        {"protocolVersion", protocol_version},
        {"capabilities", capabilities.to_json()},
        {"serverInfo", json::Value::object({
            {"name", server_info.name},
            {"version", server_info.version}
        })}
    });

    if (instructions) {
        result["instructions"] = *instructions;
    }

    return result;
}

std::optional<InitializationResult> InitializationResult::from_json(const json::Value& json) {
    if (!json.is_object()) {
        return std::nullopt;
    }

    try {
        auto protocol_version = json.get("protocolVersion", "").as_string();
        auto caps_json = json.get("capabilities", json::Value::object({}));
        auto server_info_json = json.get("serverInfo", json::Value::object({}));

        // Parse capabilities
        ServerCapabilities caps = ServerCapabilities::create_default();

        if (caps_json.contains("tools")) {
            auto tools_json = caps_json["tools"];
            if (tools_json.is_object()) {
                bool list_changed = tools_json.get("listChanged", false).as_bool();
                caps.tools = ToolsCapabilities{.list_changed = list_changed};
            }
        }

        if (caps_json.contains("resources")) {
            auto resources_json = caps_json["resources"];
            if (resources_json.is_object()) {
                bool subscribe = resources_json.get("subscribe", false).as_bool();
                bool list_changed = resources_json.get("listChanged", false).as_bool();
                caps.resources = ResourcesCapabilities{.subscribe = subscribe, .list_changed = list_changed};
            }
        }

        if (caps_json.contains("prompts")) {
            auto prompts_json = caps_json["prompts"];
            if (prompts_json.is_object()) {
                bool list_changed = prompts_json.get("listChanged", false).as_bool();
                caps.prompts = PromptsCapabilities{.list_changed = list_changed};
            }
        }

        if (caps_json.contains("logging")) {
            auto logging_json = caps_json["logging"];
            if (logging_json.is_object()) {
                std::string level_str = logging_json.get("level", "info").as_string();
                LoggingLevel level = LoggingLevel::Info;
                if (level_str == "debug") level = LoggingLevel::Debug;
                else if (level_str == "warning") level = LoggingLevel::Warning;
                else if (level_str == "error") level = LoggingLevel::Error;
                caps.logging = LoggingCapabilities{.level = level};
            }
        }

        if (caps_json.contains("progress")) {
            caps.progress = caps_json["progress"].as_bool();
        }

        // Parse server info
        types::Implementation server_info{
            .name = server_info_json.get("name", "").as_string(),
            .version = server_info_json.get("version", "").as_string()
        };

        // Parse optional instructions
        std::optional<std::string> instructions;
        if (json.contains("instructions")) {
            instructions = json["instructions"].as_string();
        }

        return InitializationResult{
            .protocol_version = std::move(protocol_version),
            .capabilities = std::move(caps),
            .server_info = std::move(server_info),
            .instructions = std::move(instructions)
        };
    } catch (...) {
        return std::nullopt;
    }
}

// ServerInfo implementation
ServerInfo::ServerInfo(std::string name, std::string version, std::optional<std::string> description)
    : name(std::move(name))
    , version(std::move(version))
    , description(std::move(description)) {}

json::Value ServerInfo::to_json() const {
    json::Value result = json::Value::object({
        {"name", name},
        {"version", version}
    });

    if (description) {
        result["description"] = *description;
    }

    return result;
}

std::optional<ServerInfo> ServerInfo::from_json(const json::Value& json) {
    if (!json.is_object()) {
        return std::nullopt;
    }

    try {
        auto name = json.get("name", "").as_string();
        auto version = json.get("version", "").as_string();

        std::optional<std::string> description;
        if (json.contains("description")) {
            description = json["description"].as_string();
        }

        return ServerInfo{
            .name = std::move(name),
            .version = std::move(version),
            .description = std::move(description)
        };
    } catch (...) {
        return std::nullopt;
    }
}

// InitializationValidator implementation
struct InitializationValidator::Impl {
    ValidationConfig config;
    mutable std::mutex mutex;

    Impl(const ValidationConfig& cfg) : config(cfg) {}
};

InitializationValidator::InitializationValidator(const ValidationConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

InitializationValidator::~InitializationValidator() = default;

void InitializationValidator::set_config(const ValidationConfig& config) {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    impl_->config = config;
}

ValidationConfig InitializationValidator::get_config() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    return impl_->config;
}

std::vector<ValidationError> InitializationValidator::validate_client_params(
    const types::InitializeRequestParams& params) const {

    std::lock_guard<std::mutex> lock(impl_->mutex);
    std::vector<ValidationError> errors;

    // Validate protocol version
    if (params.protocol_version.empty()) {
        errors.push_back(ValidationError{
            .field = "protocolVersion",
            .message = "Protocol version cannot be empty",
            .severity = ValidationSeverity::Error
        });
    } else if (impl_->config.required_protocol_version &&
               params.protocol_version != *impl_->config.required_protocol_version) {
        errors.push_back(ValidationError{
            .field = "protocolVersion",
            .message = "Unsupported protocol version: " + params.protocol_version +
                      " (required: " + *impl_->config.required_protocol_version + ")",
            .severity = ValidationSeverity::Error
        });
    }

    // Validate client info
    if (params.client_info.name.empty()) {
        errors.push_back(ValidationError{
            .field = "clientInfo.name",
            .message = "Client name cannot be empty",
            .severity = ValidationSeverity::Error
        });
    }

    if (params.client_info.version.empty()) {
        errors.push_back(ValidationError{
            .field = "clientInfo.version",
            .message = "Client version cannot be empty",
            .severity = ValidationSeverity::Warning
        });
    }

    // Validate capabilities if strict validation is enabled
    if (impl_->config.strict_capabilities) {
        const auto& caps = params.capabilities;

        // Check for required capabilities
        if (impl_->config.required_capabilities.tools && !caps.tools) {
            errors.push_back(ValidationError{
                .field = "capabilities.tools",
                .message = "Tools capability is required",
                .severity = ValidationSeverity::Error
            });
        }

        if (impl_->config.required_capabilities.resources && !caps.resources) {
            errors.push_back(ValidationError{
                .field = "capabilities.resources",
                .message = "Resources capability is required",
                .severity = ValidationSeverity::Error
            });
        }

        if (impl_->config.required_capabilities.prompts && !caps.prompts) {
            errors.push_back(ValidationError{
                .field = "capabilities.prompts",
                .message = "Prompts capability is required",
                .severity = ValidationSeverity::Error
            });
        }
    }

    return errors;
}

std::vector<ValidationError> InitializationValidator::validate_server_capabilities(
    const ServerCapabilities& caps) const {

    std::lock_guard<std::mutex> lock(impl_->mutex);
    std::vector<ValidationError> errors;

    // Validate capability combinations
    if (impl_->config.require_resources_with_tools && caps.tools && !caps.resources) {
        errors.push_back(ValidationError{
            .field = "capabilities.resources",
            .message = "Resources capability is required when tools are enabled",
            .severity = ValidationSeverity::Warning
        });
    }

    // Validate experimental capabilities
    if (caps.experimental && impl_->config.allowed_experimental_caps) {
        for (const auto& [key, value] : *caps.experimental) {
            if (impl_->config.allowed_experimental_caps->find(key) ==
                impl_->config.allowed_experimental_caps->end()) {
                errors.push_back(ValidationError{
                    .field = "capabilities.experimental." + key,
                    .message = "Experimental capability not allowed: " + key,
                    .severity = ValidationSeverity::Warning
                });
            }
        }
    }

    return errors;
}

bool InitializationValidator::is_compatible(const std::string& client_version,
                                          const std::string& server_version) const {

    std::lock_guard<std::mutex> lock(impl_->mutex);

    // Simple version compatibility check
    // In a real implementation, this might use semantic versioning
    if (impl_->config.strict_version_matching) {
        return client_version == server_version;
    }

    // Allow any server version if client doesn't specify
    if (client_version.empty()) {
        return true;
    }

    // Allow any client version if server doesn't specify
    if (server_version.empty()) {
        return true;
    }

    // Simple prefix matching for compatibility
    return client_version.substr(0, client_version.find_last_of('.')) ==
           server_version.substr(0, server_version.find_last_of('.'));
}

// Utility functions
std::string validation_severity_to_string(ValidationSeverity severity) {
    switch (severity) {
        case ValidationSeverity::Info:
            return "Info";
        case ValidationSeverity::Warning:
            return "Warning";
        case ValidationSeverity::Error:
            return "Error";
        default:
            return "Unknown";
    }
}

std::string format_validation_errors(const std::vector<ValidationError>& errors) {
    if (errors.empty()) {
        return "No validation errors";
    }

    std::stringstream ss;
    ss << "Validation errors:\n";

    for (const auto& error : errors) {
        ss << "  [" << validation_severity_to_string(error.severity) << "] "
           << error.field << ": " << error.message << "\n";
    }

    return ss.str();
}

bool has_validation_errors(const std::vector<ValidationError>& errors) {
    return std::any_of(errors.begin(), errors.end(),
        [](const ValidationError& error) {
            return error.severity == ValidationSeverity::Error;
        });
}

std::vector<ValidationError> filter_errors_by_severity(const std::vector<ValidationError>& errors,
                                                        ValidationSeverity min_severity) {
    std::vector<ValidationError> filtered;
    std::copy_if(errors.begin(), errors.end(), std::back_inserter(filtered),
        [min_severity](const ValidationError& error) {
            return error.severity >= min_severity;
        });
    return filtered;
}

// Default configurations
ValidationConfig ValidationConfig::create_default() {
    return ValidationConfig{
        .strict_capabilities = false,
        .strict_version_matching = false,
        .require_resources_with_tools = false,
        .required_capabilities = {},
        .required_protocol_version = std::nullopt,
        .allowed_experimental_caps = std::nullopt
    };
}

ValidationConfig ValidationConfig::create_strict() {
    return ValidationConfig{
        .strict_capabilities = true,
        .strict_version_matching = true,
        .require_resources_with_tools = true,
        .required_capabilities = {
            .tools = true,
            .resources = true,
            .prompts = true
        },
        .required_protocol_version = "2025-06-18",
        .allowed_experimental_caps = std::unordered_set<std::string>{"custom_capability"}
    };
}

} // namespace server
} // namespace mcp